import base64
import inspect
from functools import wraps
from inspect import isfunction

from basic.base_driver import DRIVER
from const import SCREEN_PATH


def exception_handler(func_or_class):
    # 如果是函数装饰器
    if isfunction(func_or_class):
        @wraps(func_or_class)
        def wrapper(*args, **kwargs):
            DRIVER.start_recording_screen()  # 开始录制屏幕
            try:
                return func_or_class(*args, **kwargs)
            except Exception as e:
                result = DRIVER.stop_recording_screen()
                video_data = base64.b64decode(result)
                with open(f'{SCREEN_PATH}/{func_or_class.__name__}.mp4', 'wb') as f:
                    f.write(video_data)
                raise e

        return wrapper

    # 如果是类装饰器
    elif inspect.isclass(func_or_class):
        # 遍历类的所有方法
        for name, method in inspect.getmembers(func_or_class, predicate=inspect.isfunction):
            # 只装饰以test开头的方法
            if name.startswith('test_'):
                # 应用异常处理装饰器到方法
                decorated_method = exception_handler(method)
                # 替换原方法
                setattr(func_or_class, name, decorated_method)
        return func_or_class
    else:
        return TypeError("exception_handler can only decorate functions or classes")
